#!/usr/bin/env python
"""
Key-value data extractor for various files

 Currently, this only supports columnar data files such as CSV or TSV. Other
 file formats may become supported in the future. 

 For implementation details, see below.

Author: Tom Robinson <trobinson@systemsbiology.org>
Initial Revision: February 17th, 2012

TODOs:
* Add "skipRows" to explicitly exclude rows by fixed offset or name

"""

## Credit documentation

__author__ = "Tom Robinson"
__copyright__ = "Institute for Systems Biology, 2011"
__credits__ = ["Tom Robinson"]
__license__ = "Apache 2.0"
__version__ = "1.0"
__maintainer__ = "Tom Robinson"
__email__ = "trobinson@systemsbiology.org"
__status__ = "Production"
__platform__ = "Multi-OS, Python 2.6+"


## Globals

MAX_LINE=409600000                  ## Maximum line length (in characters)
                                    ##  This, loosely, prevents DoS attacks
EMPTY_VALUES = set(["NA"])          ## Set of empty value formats


## Classes

class ColumnarKVGenerator(object):
    """ Wraps a columnar data stream or KV-alike into a generator of key-value
     pairs 

     This class provides a generator that wraps a columnar data stream passed
     in as a reader and emits a series of key-value pairs for consumption by
     downstream tools. In this way, we can extract all data elements from the
     file in a simplified manner, somewhat analogous to the Simple API for XML
     (https://en.wikipedia.org/wiki/Simple_API_for_XML), but without states.

     For further details on invocation, see __init__ and the remainder of the
     provided API. 

    """
    __reader = None         ## The reader supplied by the caller as input ##
    __keycols = None        ## The key columns to use ##
    __keycolhash = None     ## A hash table for fast column lookup ##
    __header = None         ## The header, if it exists ##
    __separator = None      ## The column separator for this stream ##
    __connector = None      ## The connector used to compose composite keys ##

    __buffer = []           ## The work buffer stored per line from reading ##
    __lineCount = 0         ## The current line index, for keyless entries ##

    def __init__(self, reader, headerCount=1, keycols=[0], separator='\t', \
      connector=':', additionalSkipLines=[0,0]): 
        """ Create a new instance of this class

         Notes: 

         * Type-correctness is currently the caller's problem. If input,
           duck-typed guarantees are not preserved, expect things to break.
        
         * The form of linebreak used when readline(...) is invoked depends on
           the open mode for the stream. Consult this article if the default
           open mode available to your platform is insufficient for your needs:
            http://www.python.org/dev/peps/pep-0278/

         * This parser assumes that headers always appear as a single line of
           the file. Complex headers, or headers that appear some distance into
           the file with other lines that should be parsed, are not currently
           supported. 

        Parameters
        ----------
        reader : streamlike
            An open reader to extract data from

        headerCount : int
            The number of header rows to parse within this reader
            Default: 1

        keycols : list of int
            The columns to use as transposed (rotated) header lines
            Default: [0]

        separator : string
            The column separator for this file
            Default: '\t'

        connector : string
            The string used to join compound keys together
            Default: ':'

        additionalSkipLines : list of int
            Additional lines to skip before and after the header
            Default: [0,0]

        Return Value
        ------------
        ColumnarKVGenerator
            A new instance of this class, configured with the given params

        Examples
        --------
        >>> myReader = open('/dev/null','r')
        >>> ckvg = ColumnarKVGenerator(myReader) # A new instance is created 

        """

        # Copy these inputs directly into internal state
        self.__reader = reader
        self.__keycols = keycols
        self.__keycolhash = dict([[i,0] for i in keycols])
        self.__separator = separator
        self.__connector = connector
        
        # If we're asked to skip additional lines before the header, do so. We
        #  still check the length, though, to curtail parsing errors.
        if additionalSkipLines[0] > 0:
            for i in xrange(0,additionalSkipLines[0]):
                line = self.__reader.readline(MAX_LINE)
                if len(line) == MAX_LINE:
                    raise IOError('Line length was too large (%s).' \
                      % str(MAX_LINE) \
                      + ' Please ensure the file is being parsed correctly.')

        # If this file has a one-line header, parse it.
        if headerCount > 0:
            for i in range(0,headerCount):
                line = self.__reader.readline(MAX_LINE)
                if len(line) == MAX_LINE:
                    raise IOError('Line length was too large (%s).' \
                      % str(MAX_LINE) \
                      + ' Please ensure the file is being parsed correctly.')
                split = line.rstrip().split(self.__separator)
                if not self.__header:
                    self.__header = split
                else:
                    # Errors in this logic will, deliberately, propagate to the
                    #  top level. Caveat emptor.
                    for idx,heading in enumerate(split):
                        self.__header[idx] += connector
                        self.__header[idx] += heading
        
        # If we're asked to skip additional lines after the header, do so. We
        #  still check the length, though, to curtail parsing errors.
        if additionalSkipLines[1] > 0:
            for i in xrange(0,additionalSkipLines[1]):
                line = self.__reader.readline(MAX_LINE)
                if len(line) == MAX_LINE:
                    raise IOError('Line length was too large (%s).' \
                      % str(MAX_LINE) \
                      + 'Please ensure the file is being parsed correctly.')

    def __iter__(self):
        """ Returns our iterator; namely, a pointer to self

        Parameters
        ----------
        None (other than the default, "self")

        Return Value
        ------------
        ColumnarKVGenerator
            A pointer back to this class

        Examples
        --------
        >>> myReader = open('/dev/null','r')
        >>> ckvg = ColumnarKVGenerator(myReader)
        >>> assert(ckvg == ckvg.__iter__()) # Points to the same thing

        """
        return self
    
    def next(self):
        """ Returns the next element, if one exists

        Parameters
        ----------
        None (other than the default, "self")

        Return Value
        ------------
        list of string
            The next element, in form [key, value]

        Examples
        --------
        >>> myReader = open('/dev/null','r')
        >>> ckvg = ColumnarKVGenerator(myReader)
        >>> try:
        ...     kv = ckvg.next() # Returns the next element, if it exists
        ... except StopIteration:
        ...     pass

        """
        
        # If our work buffer is non-empty, pop a result off the top
        if self.__buffer:
            return self.__buffer.pop()

        # Otherwise, try to read in a new line
        line = self.__reader.readline(MAX_LINE)
        if len(line) == MAX_LINE:
            raise IOError('Line length was too large (%s).' \
              % str(MAX_LINE) \
              + 'Please ensure the file is being parsed correctly.')

        self.__lineCount += 1

        # Ensure the line is non-empty (non-None should be
        # guaranteed). Otherwise, raise StopIteration().
        if line:

            splitline = line.rstrip().split(self.__separator)

            # Treat key columns as self-identifying. Any additional columnar
            # data is identified by its header, if it exists. Otherwise, it
            # will receive a number based upon its position in the file.
            rowkey = ''
            if self.__keycols:
                for keycol in self.__keycols:
                    if keycol >= len(splitline):
                        raise ValueError( \
                            'Expected key column could not be parsed: %s' % \
                             str(keycol))
                    rowkey += splitline[keycol]
                    rowkey += self.__connector
                
            

            # If no key is provided, use the current line instead
            if not rowkey:
                rowkey = str(self.__lineCount)
                rowkey += self.__connector

            # Gather up all of the columns that aren't part of the key and
            #  output them for the given row and column.
            for col in xrange(0,len(splitline)):
                if col not in self.__keycolhash:

                    # Ignore any data elements that are lexically vacuous
                    splitline[col] = splitline[col].strip()
                    if not splitline[col] or splitline[col] in EMPTY_VALUES:
                        continue
                
                    colkey = str(col)
                    if self.__header:
                        colkey = self.__header[col]

                    self.__buffer.append([rowkey + colkey,splitline[col]])

            # Finally, return the top of the buffer
            if self.__buffer:
                return self.__buffer.pop()

        # Any fall-throughs raise StopIteration()
        raise StopIteration()


# Invoke our inline doctests when this library is invoked directly
if __name__ == "__main__":
    import doctest
    doctest.testmod()
